package com.example.springbootdemo.text;

import com.example.springbootdemo.text.zh.ChineseTokenizer;
import org.wltea.analyzer.core.IKSegmenter;
import org.wltea.analyzer.core.Lexeme;

import java.io.IOException;
import java.io.StringReader;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

public class ChineseCosineSimilarity {

    /**
     * 使用 IKAnalyzer 对中文文本进行分词。
     *
     * @param text 输入的中文文本
     * @return 分词后的词语集合
     * @throws IOException 如果分词过程中出现IO错误
     */
    public static Set<String> tokenize(String text) throws IOException {
        /*Set<String> tokens = new HashSet<>();
        StringReader reader = new StringReader(text);
        IKSegmenter ik = new IKSegmenter(reader, true); // true 表示使用智能分词
        Lexeme lexeme;
        while ((lexeme = ik.next()) != null) {
            tokens.add(lexeme.getLexemeText());
        }
        reader.close();
        return tokens;*/
        return ChineseTokenizer.tokenize(text);
    }

    /**
     * 计算两个中文字符串的余弦相似度。
     *
     * @param text1 第一个中文文本字符串
     * @param text2 第二个中文文本字符串
     * @return 余弦相似度（0到1之间）
     * @throws IOException 如果分词过程中出现IO错误
     */
    public static double computeCosineSimilarity(String text1, String text2) throws IOException {
        Map<String, Integer> freq1 = getTermFrequency(text1);
        Map<String, Integer> freq2 = getTermFrequency(text2);

        // 获取所有词汇的并集
        Set<String> allTerms = new HashSet<>();
        allTerms.addAll(freq1.keySet());
        allTerms.addAll(freq2.keySet());

        // 构建向量
        double[] vector1 = new double[allTerms.size()];
        double[] vector2 = new double[allTerms.size()];
        int index = 0;
        for (String term : allTerms) {
            vector1[index] = freq1.getOrDefault(term, 0);
            vector2[index] = freq2.getOrDefault(term, 0);
            index++;
        }

        // 计算余弦相似度
        double dotProduct = 0.0;
        double normA = 0.0;
        double normB = 0.0;
        for (int i = 0; i < vector1.length; i++) {
            dotProduct += vector1[i] * vector2[i];
            normA += Math.pow(vector1[i], 2);
            normB += Math.pow(vector2[i], 2);
        }

        if (normA == 0.0 || normB == 0.0) {
            return 0.0;
        }

        return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
    }

    /**
     * 计算文本的词频。
     *
     * @param text 输入的中文文本
     * @return 词频映射
     * @throws IOException 如果分词过程中出现IO错误
     */
    private static Map<String, Integer> getTermFrequency(String text) throws IOException {
        Set<String> tokens = tokenize(text);
        Map<String, Integer> freqMap = new HashMap<>();
        for (String token : tokens) {
            freqMap.put(token, freqMap.getOrDefault(token, 0) + 1);
        }
        return freqMap;
    }

    public static void main(String[] args) {
        String text1 = "我喜欢学习人工智能和机器学习。";
        String text2 = "机器学习和人工智能是我的兴趣所在。";

        try {
            double similarity = computeCosineSimilarity(text1, text2);
            System.out.printf("余弦相似度: %.4f\n", similarity);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}
